__author__ = 'sdjkx'

# 扫描XML文件，把中文标签替换为英文标签
import os
from libs.ustr import *

try:
  import xml.etree.cElementTree as ET
except ImportError:
  import xml.etree.ElementTree as ET
from PyQt5.QtWidgets import QMessageBox,QFileDialog
import codecs
from xml.etree.ElementTree import Element, SubElement
from lxml import etree

from libs.constants import DEFAULT_ENCODING
ENCODE_METHOD = DEFAULT_ENCODING

class ScanXMLReplaceLabel:

    def __init__(self, labelImg, label_dict):
        self.labelImg = labelImg
        self.label_dict = label_dict

    def scanCPLabelFun(self):
        ret = QMessageBox.question(self.labelImg, '提示', "确定要执行替换XML文件中文标签？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if ret == QMessageBox.No:
            return 

        path = '.'
        dirpath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择要替换中文标签的XML文件路径", path,  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))

        if dirpath is not None and len(dirpath) > 1:
            ret = QMessageBox.question(self.labelImg, '提示', "确定要扫描目录【"+dirpath+"】里的XML文件吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
            if ret == QMessageBox.Yes:
                self.scanCpAndBusLabel(dirpath)
                return

    def scanCpAndBusLabel(self, xmlPath):
        xmlFileList = os.listdir(xmlPath)
        for xmlFileName in xmlFileList:
            xmlFilePath =  os.path.join(xmlPath, xmlFileName)

            if os.path.exists(xmlFilePath):
                print("XML文件：" + xmlFilePath)
                if not xmlFilePath.endswith("xml"):
                    continue

                tree = ET.parse(xmlFilePath)     # 打开xml文档
                root = tree.getroot()         #获得root节点

                obj_hadReplace =  root.find("hadReplace")
                if obj_hadReplace != None:
                    continue

                for object in root.findall('object'): #找到root节点下的所有object节点
                    obj = object.find('name')
                    label_name = obj.text   # 子节点下节点name的值

                    labelObj = self.label_dict[label_name]
                    if labelObj != None:
                        obj.text = labelObj.get_labelEn()
                    else:
                        print('XML文件 ' + xmlFilePath + ' 中的标签' + label_name + '不存在')

                hadReplace = SubElement(root, 'hadReplace')
                hadReplace.text = '1'
                self.save(root, xmlFilePath)
            else:
                print('XML文件不存在:' + xmlFilePath)

        QMessageBox.information(self.labelImg, u'提示', "XML文件标签替换完成")
            
    def prettify(self, elem):
        """
            Return a pretty-printed XML string for the Element.
        """
        rough_string = ET.tostring(elem, 'utf8')
        root = etree.fromstring(rough_string)
        return etree.tostring(root, pretty_print=True, encoding=ENCODE_METHOD).replace("  ".encode(), "\t".encode())
        # minidom does not support UTF-8

    def save(self,root, targetFile=None):
        out_file = codecs.open(targetFile, 'w', encoding=ENCODE_METHOD)

        prettifyResult = self.prettify(root)
        out_file.write(prettifyResult.decode('utf8'))
        out_file.close()